import torch
from torch import nn
from torch import optim
from torch.nn import functional as F

from torchvision import datasets
from torchvision import transforms

from matplotlib import pyplot
import numpy as np
import os.path

model_save_path = os.path.realpath('pytorch/jizhi/figure_identify/figure_identify_cnn_model')
print(f'model save path:{model_save_path}')

image_size = 28
num_classes = 10
num_epochs = 1
batch_size = 64

def rightness(predictions, labels):
    pred = torch.max(predictions.data, 1)[1]
    rights = pred.eq(labels.data.view_as(pred)).sum()
    return rights, len(labels)

class CNNUtil():
    def __init__(self):
        self.prepare_data()
        #self.test_show()
        
    def prepare_data(self):
        print('begin prepare data')
        self.train_data = datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor(), download=True)
        self.test_data = datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())
        
        indices = range(len(self.test_data))
        self.indices_verify = indices[:5000]
        self.indices_test = indices[5000:]     
        
        self.sampler_verify = torch.utils.data.sampler.SubsetRandomSampler(self.indices_verify)
        self.sampler_test = torch.utils.data.sampler.SubsetRandomSampler(self.indices_test)
        
        self.train_loader = torch.utils.data.DataLoader(dataset=self.train_data, batch_size=batch_size, shuffle=True)
        self.validation_loader = torch.utils.data.DataLoader(dataset=self.test_data, batch_size=batch_size, shuffle=False, sampler=self.sampler_verify)
        self.test_loader = torch.utils.data.DataLoader(dataset=self.test_data, batch_size=batch_size, shuffle=False, sampler=self.sampler_test)

    def test_show(self):
        print('begin test show')
        index = 200
        img_np = self.train_data[index][0].numpy()
        pyplot.imshow(img_np[0, ...])
        pyplot.show()
        
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        self.depth = [4, 8] # number of conv core in every layer
        self.conv1 = nn.Conv2d(1, 4, 5, padding=2)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(self.depth[0], self.depth[1], 5, padding=2)
        self.fc1 = nn.Linear((image_size // 4)**2 * self.depth[1], 512)
        self.fc2 = nn.Linear(512, num_classes)
        
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        
        x = x.view(-1, (image_size // 4)**2 * self.depth[1])
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        x = F.log_softmax(x, dim=1)
        return x
    
    def retrieve_features(self, x):
        feature_map1 = F.relu(self.conv1(x))
        x = self.pool(feature_map1)
        feature_map2 = F.relu(self.conv2(x))
        return feature_map1, feature_map2

if __name__ == '__main__':
    cnn_util = CNNUtil()
    cnn = CNN()
    
    cost = nn.CrossEntropyLoss()
    optimizer = optim.SGD(cnn.parameters(), lr=0.01, momentum=0.9)
    
    record, weights = [], []
    print('begin train')
    for epoch in range(num_epochs):
        train_rights = []
        for batch_index, (data, target) in enumerate(cnn_util.train_loader):
            data, target = data.clone().requires_grad_(True), target.clone().detach()
            cnn.train()
            
            output = cnn(data)
            loss = cost(output, target)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            right = rightness(output, target)
            train_rights.append(right)
            if batch_index % 200 == 0:
                cnn.eval()
                val_rights = []
                for (data, target) in cnn_util.validation_loader:
                    data, target = data.clone().requires_grad_(True), target.clone().detach()
                    output = cnn(data)
                    right = rightness(output, target)
                    val_rights.append(right)
                train_r = (sum([tup[0] for tup in train_rights]), sum([tup[1] for tup in train_rights]))
                val_r = (sum([tup[0] for tup in val_rights]), sum([tup[1] for tup in val_rights]))
                print(f'no.{epoch}, {batch_index*batch_size}/{len(cnn_util.train_loader.dataset)}')
                print(f'loss:{loss.data}, train accu:{train_r[0]/train_r[1]}, verify accu:{val_r[0]/val_r[1]}')
                record.append((100 - 100 * train_r[0] / train_r[1], 100 - 100 * val_r[0] / val_r[1]))
                weights.append([cnn.conv1.weight.data.clone(), cnn.conv1.bias.data.clone(), cnn.conv2.weight.data.clone(), cnn.conv2.bias.data.clone()])
    cnn.eval()
    vals = []
    for data, target in cnn_util.test_loader:
        data, target = data.clone().detach().requires_grad_(True), target.clone().detach()
        output = cnn(data)
        val = rightness(output, target)
        vals.append(val)
    rights = (sum([tup[0] for tup in vals]), sum([tup[1] for tup in vals]))
    right_rate = rights[0] / rights[1]
    print(f'test scores:{right_rate}')
    # save cnn model
    # torch.save(cnn, model_save_path)
    # plot
    '''
    pyplot.figure(figsize=(10, 7))
    pyplot.plot(record)
    pyplot.xlabel('steps')
    pyplot.ylabel('error rate')
    pyplot.show()
    '''
    # plot no.1 conv cores
    '''
    print('plot conv1 cores')
    pyplot.figure(figsize=(8, 6))
    for i in range(4):
        pyplot.subplot(1, 4, i + 1)
        pyplot.imshow(cnn.conv1.weight.data.numpy()[i, 0, ...])
    pyplot.show()
    # '''
    # plot no.1 conv feature map
    index = 1000
    input_x = cnn_util.test_data[index][0].unsqueeze(0)
    feature_maps = cnn.retrieve_features(input_x)
    '''
    index = 1000
    input_x = cnn_util.test_data[index][0].unsqueeze(0)
    feature_maps = cnn.retrieve_features(input_x)
    pyplot.figure(figsize=(8, 6))
    for i in range(4):
        pyplot.subplot(1, 4, i + 1)
        pyplot.imshow(feature_maps[0][0, i, ...].data.numpy())
    pyplot.show()
    '''
    '''
    # plot conv2 core
    pyplot.figure(figsize=(15, 10))
    for i in range(4):
        for j in range(8):
            pyplot.subplot(4, 8, i * 8 + j + 1)
            pyplot.imshow(cnn.conv2.weight.data.numpy()[j, i, ...])
    pyplot.show()
    '''
    '''
    plot conv2 feature map
    pyplot.figure(figsize=(8, 6))
    for i in range(8):
        pyplot.subplot(2, 4, i + 1)
        pyplot.imshow(feature_maps[1][0, i, ...].data.numpy())
    pyplot.show()
    '''
    